﻿using System;
using System.Runtime.Remoting.Contexts;
using System.Runtime.Remoting.Messaging;
using System.Transactions;

namespace ConsoleApplication5
{
    [AttributeUsage(AttributeTargets.Method, AllowMultiple = false, Inherited = false)]
    public class Trans : ContextAttribute, IContributeObjectSink
    {
        public Trans(string name) : base(name)
        {
        }

        public IMessageSink GetObjectSink(MarshalByRefObject obj, IMessageSink nextSink)
        {
            return new AopTrans(nextSink);
        }

        class AopTrans : IMessageSink
        {
            public AopTrans(IMessageSink nextSink)
            {
                NextSink = nextSink;
            }

            public IMessageSink NextSink { get; private set; }

            public IMessage SyncProcessMessage(IMessage msg)
            {
                Preprocess();
                IMessage retMsg = NextSink.SyncProcessMessage(msg);
                Postprocess();
                return retMsg;
            }

            public IMessageCtrl AsyncProcessMessage(IMessage msg, IMessageSink replySink)
            {
                throw new NotImplementedException();
            }

            private void Preprocess()
            {
                TransactionScope ts = new TransactionScope(TransactionScopeOption.Required);
                CallContext.SetData(ContextName, ts);
            }

            private void Postprocess()
            {
                TransactionScope ts = CallContext.GetData(ContextName) as TransactionScope;
                if (ts != null)
                {
                    ts.Complete();
                    ts.Dispose();
                }
            }

            public static readonly string ContextName = "AopTrans";
        }
    }
}
